import torch
from torch import nn
from torchsummary import summary
import os

from dataset import GetDataLoader

## 残差块
class Residual(nn.Module):
    def __init__(self,in_channels,out_channels,stride=1,use_1conv=False):
        '''
        :param in_channels: 输入的通道
        :param out_channels: 输出的通道
        :param stride: 第一个卷积的步幅
        :param use_1conv: 另一个路是否有1x1卷积
        卷积1：in-->out,卷积2：out-->out
        '''
        super(Residual,self).__init__()

        self.Relu = nn.ReLU()
        # 残差块的第1个卷积，步幅可能是1或2
        self.conv1 = nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=3,stride=stride,padding=1)
        # 残差块的第2个卷积，步幅固定是1
        self.conv2 = nn.Conv2d(in_channels=out_channels,out_channels=out_channels,kernel_size=3,stride=1,padding=1)

        self.bn1 = nn.BatchNorm2d(num_features=out_channels)
        self.bn2 = nn.BatchNorm2d(num_features=out_channels)
        if use_1conv:
            # 残差块中1x1卷积步幅固定是2
            self.conv3 = nn.Conv2d(in_channels=in_channels,out_channels=out_channels,kernel_size=1,stride=2,padding=0)
        else:
            self.conv3 = None

    def forward(self,x):
        y = self.Relu(self.bn1(self.conv1(x)))
        y = self.bn2(self.conv2(y))

        if self.conv3:
            x = self.conv3(x)
        y = self.Relu(y+x)
        return y


## resnet18网络
class ResNet18(nn.Module):
    # 图片通道
    IMG_CHANNELS = GetDataLoader().IMG_CHANNELS
    # 图片分类数
    IMG_CLASSES = len(GetDataLoader().IMG_CLASSES)
    # 模型参数文件保存位置
    PTH_PATH = os.path.join('./model_files','best_model.pth')
    def __init__(self):

        super(ResNet18,self).__init__()
        # 第一块
        self.b1 = nn.Sequential(
            nn.Conv2d(in_channels=self.IMG_CHANNELS,out_channels=64,kernel_size=7,stride=2,padding=3),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3,stride=2,padding=1)
        )
        # 第2块：残差块1和2都是无1x1卷积
        self.b2 = nn.Sequential(
            Residual(in_channels=64,out_channels=64,stride=1,use_1conv=False),
            Residual(in_channels=64, out_channels=64,stride=1,use_1conv=False),
        )

        # 第3块：残差块1有1x1卷积，残差块2无1x1卷积
        self.b3 = nn.Sequential(
            Residual(in_channels=64, out_channels=128, stride=2, use_1conv=True),
            Residual(in_channels=128, out_channels=128, stride=1, use_1conv=False),
        )

        # 第4块：残差块1有1x1卷积，残差块2无1x1卷积
        self.b4 = nn.Sequential(
            Residual(in_channels=128, out_channels=256, stride=2, use_1conv=True),
            Residual(in_channels=256, out_channels=256, stride=1, use_1conv=False),
        )

        # 第5块：残差块1有1x1卷积，残差块2无1x1卷积
        self.b5 = nn.Sequential(
            Residual(in_channels=256, out_channels=512, stride=2, use_1conv=True),
            Residual(in_channels=512, out_channels=512, stride=1, use_1conv=False),
        )

        # 第6块
        self.b6 = nn.Sequential(
            # 全局平均池化
            nn.AdaptiveAvgPool2d((1,1)),
            #展平层
            nn.Flatten(),
            # 全连接层（参数多少就多少个神经元）
            nn.Linear(in_features=512,out_features=self.IMG_CLASSES)
        )

    def forward(self,x):
        x = self.b1(x)
        x = self.b2(x)
        x = self.b3(x)
        x = self.b4(x)
        x = self.b5(x)
        x = self.b6(x)
        return x

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = ResNet18().to(device)
    print(summary(model,(1,224,224)))